import collections
import os

from PIL import Image
from tqdm import tqdm

from option import parse_args

MAX_ACTOR = 5


def count_actors(video_path):
    max_factor = 0
    for crop_name in os.listdir(video_path):
        crop_name_split = crop_name[:-4].split('_')
        assert len(crop_name_split) == 2
        max_factor = max(max_factor, int(crop_name_split[1]))
    return max_factor


def get_frame_min_max_width(video_path, frame):
    w_min, w_max = 500, 0
    for i in range(MAX_ACTOR):
        image_path = os.path.join(video_path, "{}_{}.png".format(frame, i))
        if os.path.exists(image_path):
            with Image.open(image_path) as img:
                w_min = min(w_min, img.size[0])
                w_max = max(w_max, img.size[0])
    return w_min, w_max


def multiple_actors(original_crops, fake_crops, func, template):
    # real
    video_names = os.listdir(original_crops)
    video_names.sort()
    for video_name in tqdm(video_names):
        video_path = os.path.join(original_crops, video_name)
        num = 1
        dd = collections.defaultdict(int)
        for crop_name in os.listdir(video_path):
            crop_name_split = crop_name[:-4].split('_')
            assert len(crop_name_split) == 2
            dd[int(crop_name_split[0])] = max(dd[int(crop_name_split[0])], int(crop_name_split[1]) + 1)
        for k in dd.keys():
            num = max(num, dd[k])
        if num > 1:
            func(template.format(video_name, num))

    # fake
    video_names = os.listdir(fake_crops)
    video_names.sort()
    for video_name in tqdm(video_names):
        video_path = os.path.join(fake_crops, video_name)
        num = 1
        dd = collections.defaultdict(int)
        for crop_name in os.listdir(video_path):
            crop_name_split = crop_name[:-4].split('_')
            assert len(crop_name_split) == 2
            dd[int(crop_name_split[0])] = max(dd[int(crop_name_split[0])], int(crop_name_split[1]) + 1)
        for k in dd.keys():
            num = max(num, dd[k])
        if num > 1:
            func(template.format(video_name, num))


def main():
    args = parse_args()
    version = args.compression_version
    original_crops = os.path.join(args.root_dir, 'original_sequences', 'actors', version, 'crops')
    fake_crops = os.path.join(args.root_dir, 'manipulated_sequences', 'DeepFakeDetection', version, 'crops')
    args.save = True
    if args.save:
        with open(f'multiple_actors_statistics_{version}.dat', 'w') as f:
            multiple_actors(original_crops, fake_crops, f.write, '{},{}\n')
    else:
        multiple_actors(original_crops, fake_crops, print, '{},{}')


if __name__ == '__main__':
    main()
